from BackTestingRunner import BackTestingRunner
from Strategy import Strategy


class MaGridStrategy(Strategy):
    def __init__(self):
        super().__init__()
        self.context = None
        self.stopPrice = 0 # 止损价
        self.openPrice = 0 # 开仓价格
        self.openPos = 0.5 # 开仓仓位

    def onKlineStep(self, index, context):
        """
        上证：价格突破 ma20 则买入， atr 下边缘为止损点， 做网格交易
             3000 以下买入的不设止损
        """
        self.context = context
        # 昨日价格
        prePrice = context.getNewPrice(index-1)
        # 今日价格
        newPrice = context.getNewPrice(index)

        ma20 = context.getMa20(index)

        if prePrice < ma20 < newPrice:
            # 价格突破 20 均线
            self.handleBuy(newPrice, index)
        self.handleStop(newPrice, index)
        self.handleGrid(newPrice, index)

    def handleBuy(self, newPrice, index):
        # 判断有无持仓
        if self.repo.getKeepNum() <= 0:
            self.repo.buyByRate(self.openPos,index) # 开仓 3 层
            self.openPrice = newPrice
            if newPrice < 3000:
                # 3000 点以下不设止损
                self.stopPrice = 0
            else:
                self.stopPrice = self.context.getAtrLow(index) # 设置止损价格
            print(f'开仓, 成本价={newPrice}, 止损价格: {self.stopPrice}')

    def handleGrid(self, newPrice, index):
        if self.repo.getKeepNum() <= 0:
            return
        """ 执行网格加仓或减仓
           开仓 0.3 仓位，开仓价格往止损方向越跌越加仓，
           划分 7 个网格， 每跌一格加仓一层 0.1
         """
        if newPrice < 3000 and self.stopPrice == 0:
            keepPos = self.repo.getKeepPos()
            if keepPos > 0.99:
                return
            # 3000 点以下，定投, 每下降 1% 加仓 1 层
            backRate = round((3000 - newPrice) / 3000, 1) * 100
            needPos = self.openPos + round(backRate / 1, 0) * 0.1
            if needPos > keepPos + 0.1:
                self.repo.buyByRate(needPos, index)
            return
        if self.stopPrice == 0:
            # 3000 点以下开仓的， 到 3000 以上， 切价格下破 ma20 止盈
            ma20 = self.context.getMa20(index)
            if newPrice < ma20:
                print(f'止盈清仓：newPrice = {newPrice}')

                self.repo.sellByRate(1, index)
            return
        # 网格空间
        gridH = (self.openPrice - self.stopPrice)
        # 回撤高度
        backH = (self.openPrice - newPrice)
        if backH > 0:
            needKeepPos = self.openPos + round(backH/gridH, 1)
            keepPos = self.repo.getKeepPos()
            if needKeepPos < 0.1:
                needKeepPos = 0
            if needKeepPos > keepPos + 0.1:
                print(f'加仓：needPos = {needKeepPos},gridH = {gridH}, backH = {backH}')
                self.repo.buyByRate(needKeepPos, index)
            elif needKeepPos < keepPos - 0.1:
                print(f'减仓：needPos = {needKeepPos},gridH = {gridH}, backH = {backH}')
                self.repo.sellByRate(keepPos - needKeepPos, index)
        if backH + gridH <= 0:
            print(f'止盈清仓：newPrice = {newPrice},gridH = {gridH}, backH = {backH}')

            self.repo.sellByRate(1, index)

    def handleStop(self, newPrice, index):
        """止损逻辑"""
        if self.repo.getKeepNum() and self.stopPrice > newPrice:
            self.repo.sellByRate(1, index)
            print(f'止损清仓, 成本价={newPrice}, 止损价格: {self.stopPrice}')


BackTestingRunner('sh000001', 240 * 20, '2025-3-01').start(MaGridStrategy(), 5000000)
